package tarantula.instrumenter;

import java.util.Iterator;

import java.util.Map;

import org.apache.bcel.generic.Type;

import plume.Option;

import soot.Body;
import soot.BodyTransformer;
import soot.Local;
import soot.Modifier;
import soot.RefType;
import soot.Scene;
import soot.SootClass;
import soot.SootField;
import soot.SootMethod;
import soot.Trap;
import soot.Unit;
import soot.Value;
import soot.VoidType;
import soot.jimple.CaughtExceptionRef;
import soot.jimple.Expr;
import soot.jimple.IdentityStmt;
import soot.jimple.IntConstant;
import soot.jimple.Jimple;
import soot.jimple.NullConstant;
import soot.jimple.StringConstant;
import soot.util.Chain;
import tarantula.util.Globals;
import soot.BooleanType;

public class JUnitTestInstrumenter extends BodyTransformer {
	
	@Option("The output dir of all trace files that record coverage information")
	public static String trace_file_dir = "." + Globals.fileSep + "testdata" + Globals.fileSep + "tracedir";

	public static final String JUNIT_TEST_CASE = "junit.framework.TestCase";
	
	public static final String[] SKIPPED_PACKAGES = {"junit.", /**"java.",*/ "javax.", "sun."};
	
	@Override
	protected void internalTransform(Body body, String phase, Map options) {
		SootMethod method = body.getMethod();
		SootClass clazz = method.getDeclaringClass();
		//skip all system class
		if(shouldSkipThisClass(clazz)) {
			return;
		}
		if(!isJUnitTestMethod(method)) {
			return;
		}
		
		//construct the unique output trace file name for each JUnit test
		String clazzName = clazz.getName();
		String methodName = method.getName();
		String outputFilePath = trace_file_dir + Globals.fileSep +
		    clazzName + "." + methodName + Tracer.trace_file_suffix;
		
		//start transformation below
		//See comments above Tracer#flushAndClean a transformation process
		
		/******************************************************************/
		/** All instrumentation code goes here **/
		
	/** original test:
	 * 
	 * public void testMethodName() {
	 *     code_block
	 * }
	 * 
	 * after transformation:
	 * 
	 * public void testMethodName() *throws Throwable* {   //change the signature
	 *     Tracer.startRecordExecution();   //this must be executed
	 *     boolean isFailed = false;
	 *     try {
	 *         code_block;
	 *     } catch (Throwable e) {
	 *         isFailed = true;      //the test fails
	 *         throw e;
	 *     } finally {
	 *         Tracer.flushAndClean(outputFilePath,  isFailed); //outputFilePath is defined above
	 *     }
	 * }
	 */
		
	/**
	 * (commented by Cheng Zhang) Actually we do the transformation as below:
	 * 
	 * original test:
	 * 
	 * public void testMethodName() {
	 *     code_block
	 * }
	 * 
	 * after transformation:
	 * 
	 * public void testMethodName() {
	 *     Tracer.startRecordExecution();   //this must be executed
	 *     Throwable ee = null;
	 *     try {
	 *         code_block;
	 *     } catch (Throwable e) {
	 *         ee = e;
	 *     } 
	 *     if(ee == null){
	 *       Tracer.cleanAndFlush(false); // false means the test does not fail
	 *     } else {
	 *     	 Tracer.cleanAndFlush(true); // true means the test fails
	 *       throw ee;
	 *     }
	 * }
	 */
	
		Chain<Unit> unitChain = body.getUnits();
		
		Unit firstUnit = unitChain.getFirst();
		Unit lastUnit = unitChain.getLast();
		
		CaughtExceptionRef exceptionRef = Jimple.v().newCaughtExceptionRef();
		Local exceptionLocal = Jimple.v().newLocal("_exception_Local", RefType.v("java.lang.Throwable"));
		body.getLocals().add(exceptionLocal);
		Unit exceptionHandler = Jimple.v().newIdentityStmt(exceptionLocal, exceptionRef);
		
		Unit extraReturnStmt = Jimple.v().newReturnVoidStmt();
		unitChain.add(exceptionHandler);
		unitChain.add(extraReturnStmt);
		
		Local exceptionLocal2 = Jimple.v().newLocal("_exception_Local_02", RefType.v("java.lang.Throwable"));
		body.getLocals().add(exceptionLocal2);
		Unit exceptionAssign= Jimple.v().newAssignStmt(exceptionLocal2, exceptionLocal);
		unitChain.insertAfter(exceptionAssign, exceptionHandler);
		
		Unit nullAssign = Jimple.v().newAssignStmt(exceptionLocal2, NullConstant.v());
		Unit firstNonIdentityStmt = null;
		Iterator<Unit> unitIte = unitChain.iterator();
		while(unitIte.hasNext()){
			Unit tmpUnit = unitIte.next();
			if(!(tmpUnit instanceof IdentityStmt)){
				firstNonIdentityStmt = tmpUnit;
				break;
			}
		}
		if(firstNonIdentityStmt != null){
			unitChain.insertBefore(nullAssign, firstNonIdentityStmt);
			
			SootMethod startRecordMethod = Scene.v().getMethod("<tarantula.instrumenter.Tracer: void startRecordExecution()>");
			Expr startRecordExpr = Jimple.v().newStaticInvokeExpr(startRecordMethod.makeRef());
			Unit startRecordStmt = Jimple.v().newInvokeStmt(startRecordExpr);
			unitChain.insertBefore(startRecordStmt, nullAssign);
		}
		
		SootClass throwableClass = Scene.v().getSootClass("java.lang.Throwable");
		Trap trap = Jimple.v().newTrap(throwableClass, firstUnit, exceptionHandler, exceptionHandler);
		Chain<Trap> trapChain = body.getTraps();
		trapChain.add(trap);
		
		// insert a call to Tracer.flushAndClean() to record a test failure
		Local failureTrueLocal = Jimple.v().newLocal("_failure_true_local", BooleanType.v());
		body.getLocals().add(failureTrueLocal);
		Unit failureTrueAssign = Jimple.v().newAssignStmt(failureTrueLocal, IntConstant.v(1));
		unitChain.insertAfter(failureTrueAssign, exceptionAssign);
		
		SootClass tracerClass = Scene.v().getSootClass("tarantula.instrumenter.Tracer");
		SootMethod flushAndCleanMethod = Scene.v().getMethod("<tarantula.instrumenter.Tracer: void flushAndClean(java.lang.String,boolean)>");
		Expr flushAndCleanExpr = Jimple.v().newStaticInvokeExpr(flushAndCleanMethod.makeRef(), StringConstant.v(outputFilePath), failureTrueLocal);
		Unit flushAndCleanStmt = Jimple.v().newInvokeStmt(flushAndCleanExpr);
		unitChain.insertAfter(flushAndCleanStmt, failureTrueAssign);
		
		// insert a throw statement
		Unit reThrowStmt = Jimple.v().newThrowStmt(exceptionLocal2);
		unitChain.insertAfter(reThrowStmt, flushAndCleanStmt);
		
		// insert a call to Tracer.flushAndClean() to record a test failure
		Local failureFalseLocal = Jimple.v().newLocal("_failure_false_local", BooleanType.v());
		body.getLocals().add(failureFalseLocal);
		Unit failureFlaseAssign = Jimple.v().newAssignStmt(failureFalseLocal, IntConstant.v(0));
		unitChain.insertAfter(failureFlaseAssign, reThrowStmt);
		
		Expr flushAndCleanExpr2 = Jimple.v().newStaticInvokeExpr(flushAndCleanMethod.makeRef(), StringConstant.v(outputFilePath), failureFalseLocal);
		Unit flushAndCleanStmt2 = Jimple.v().newInvokeStmt(flushAndCleanExpr2);
		unitChain.insertAfter(flushAndCleanStmt2, failureFlaseAssign);
		
		// insert an if statement
		Value nullEq = Jimple.v().newEqExpr(exceptionLocal2, NullConstant.v());
		Unit ifStmt = Jimple.v().newIfStmt(nullEq, failureFlaseAssign);
		unitChain.insertBefore(ifStmt, failureTrueAssign);
		
		// insert a goto statement
		Unit gotoStmt = Jimple.v().newGotoStmt(ifStmt);
		unitChain.insertBefore(gotoStmt, lastUnit);
	}
	
	private boolean shouldSkipThisClass(SootClass clazz) {
		String pkgName = clazz.getPackageName();
		for(String skipped : SKIPPED_PACKAGES) {
			if(pkgName.startsWith(skipped)) {
				return true;
			}
		}
		return false;
	}

	private boolean isJUnitTestMethod(SootMethod method) {
		String methodName = method.getName();
		//the name should start with "test"
		if(!methodName.startsWith("test")) {
			return false;
		}
		//the signature should be public void, and non-static, and no parameters
		if(!Modifier.isPublic(method.getModifiers()) || Modifier.isStatic(method.getModifiers()) 
				|| !method.getReturnType().getClass().equals(VoidType.class)
				|| method.getParameterCount() != 0) {
			return false;
		}
		//check if this class extends junit.framework.TestCase
		// FIXME: it seems that in JUnit 4, a test case does not have to extend junit.framework.TestCase
		boolean extendTestCase = false;
		SootClass clazz = method.getDeclaringClass();
		SootClass superClass = clazz.getSuperclass();
		while(superClass != null) {
			if(superClass.getName().equals(JUNIT_TEST_CASE)) {
				extendTestCase = true;
				break;
			}else if(superClass.getName().equals("java.lang.Object")){
				break;
			}
			superClass = superClass.getSuperclass();
		}
		
		return extendTestCase;
	}
}
